import cv2
import numpy as np

def resize_images(img1, img2):
    """调整两幅图像的尺寸，使其具有相同的最大宽度和高度"""
    # 获取图像的尺寸
    h1, w1 = img1.shape[:2]
    h2, w2 = img2.shape[:2]
    
    # 计算目标尺寸
    max_w = max(w1, w2)
    max_h = max(h1, h2)
    
    # 调整图像尺寸
    img1_resized = cv2.resize(img1, (max_w, max_h))
    img2_resized = cv2.resize(img2, (max_w, max_h))
    
    return img1_resized, img2_resized

def align_images(ref_image, target_image):
    """对齐两幅图像"""
    # 创建 ORB 特征检测器
    orb = cv2.ORB_create(nfeatures=1000)
    
    # 检测特征点和描述子
    kp1, des1 = orb.detectAndCompute(ref_image, None)
    kp2, des2 = orb.detectAndCompute(target_image, None)
    
    # 使用暴力匹配器匹配特征点
    bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
    matches = bf.match(des1, des2)
    
    # 按匹配质量排序
    matches = sorted(matches, key=lambda x: x.distance)
    
    # 选择前 70% 的匹配点
    num_matches = int(len(matches) * 0.7)
    matches = matches[:num_matches]
    
    # 提取匹配的特征点
    src_pts = np.float32([kp1[m.queryIdx].pt for m in matches]).reshape(-1, 2)
    dst_pts = np.float32([kp2[m.trainIdx].pt for m in matches]).reshape(-1, 2)
    
    # 计算单应矩阵
    M, mask = cv2.findHomography(dst_pts, src_pts, cv2.RANSAC, 5.0)
    
    # 对齐目标图像
    aligned_image = cv2.warpPerspective(target_image, M, (ref_image.shape[1], ref_image.shape[0]))
    
    return aligned_image

def preprocess_images(img1, img2):
    """预处理两幅图像"""
    # 调整图像尺寸
    img1, img2 = resize_images(img1, img2)
    
    # 对齐图像
    aligned_image = align_images(img1, img2)
    
    # 转换为灰度图
    img1_gray = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    img2_gray = cv2.cvtColor(aligned_image, cv2.COLOR_BGR2GRAY)
    return img1_gray, img2_gray
    
    # # 归一化图像
    # img1_norm = (img1_gray - np.mean(img1_gray)) / np.std(img1_gray)
    # img2_norm = (img2_gray - np.mean(img2_gray)) / np.std(img2_gray)
    
    # return img1_norm, img2_norm

# 读取图像
img1 = cv2.imread('image1.jpg')
img2 = cv2.imread('draw_2.jpg')

# 预处理图像
img1_processed, img2_processed = preprocess_images(img1, img2)

cv2.imwrite('img1_processed.jpg', img1_processed)
cv2.imwrite('img2_processed.jpg', img2_processed)

# # 显示结果
cv2.imshow('Image 1 Processed', img1_processed)
cv2.imshow('Image 2 Processed', img2_processed)
cv2.waitKey(0)
cv2.destroyAllWindows()